# Author: 空间栈
# Gitee: https://gitee.com/shujianyang/gras-data-processing

import argparse
from pathlib import Path
import re

from pds4_tools import pds4_read       # to read and inspect the data and metadata
import numpy as np
from PIL import Image

from skimage import exposure
from skimage import data, img_as_float
import colour
from colour.plotting import *
from tqdm import tqdm

from colour_demosaicing import (
    demosaicing_CFA_Bayer_bilinear,
    demosaicing_CFA_Bayer_Malvar2004,
    demosaicing_CFA_Bayer_Menon2007,
    mosaicing_CFA_Bayer)
cctf_encoding = colour.cctf_encoding
colour.utilities.filter_warnings()

def read_pds(path):
    data = pds4_read(path, quiet=True)
    img = np.array(data[0].data)
    img = img_as_float(img)
    return img

def debayer_img(img, CFA='RGGB'):
    # Menon2007 yields better edges than bilinear
    debayered = cctf_encoding(demosaicing_CFA_Bayer_Menon2007(img, CFA))
    return debayered

def stretch_img(img):
    p2, p98 = np.percentile(img, (2, 98))
    img = exposure.rescale_intensity(img, in_range=(p2, p98))
    return img

def export_img(name, img):
    pil_img = Image.fromarray(np.uint8(img*255))
    pil_img.save(name)

#根据文件名判断图像是否为彩色而不是单色
#单色图像主要出现在嫦娥三、四号PCAM拍摄的照片中，表现为在载荷名后有-Q的后缀
#Find out if an image is a color image (relative to panchromatic image).
#Panchromatic images are usually found in CE-3 and CE-4 PCAM photos,
# indicated by a "-Q" suffix of the payload parameter.
def is_color_image(file):
    name = file.name
    return re.search('PCAM[LR]-Q', name) == None

ap = argparse.ArgumentParser()
ap.add_argument('-i', '--inputdir', required=True,
                help='Input folder where data files are stored.')
ap.add_argument('-o', '--outputdir', 
                help='Output folder where images are exported.')
ap.add_argument('-f', '--format', choices=['jpg', 'png'], default='jpg',
                help='Output image format. Can be jpg or png.')
# ap.add_argument('--debayer', action=argparse.BooleanOptionalAction, default=True,
#                 help='Enable or disable debayer process.')
# ap.add_argument('--stretch', action=argparse.BooleanOptionalAction, default=True,
#                 help='Enable or disable stretch process.')
args = vars(ap.parse_args())

dir_path = args['inputdir']
output_ext = args['format']
file_list = list(Path(dir_path).glob('*2BL'))
file_list.extend(Path(dir_path).glob('*2CL'))

for file in tqdm(file_list):
    try:
        img = read_pds(str(file))
        if file.suffix == '.2BL':
            if args['debayer'] and is_color_image(file):
                img = debayer_img(img)
            if args['stretch']:
                img = stretch_img(img)
        if args['outputdir']:
            export_img(f"{Path(args['outputdir']).joinpath(file.stem)}.{output_ext}", img)
        else:
            export_img(f"{str(file)}.{output_ext}", img)
    except:
        print(str(file))